/*
 * Copyright (c) China Telecom Cloud Technology Co., Ltd. 2024-2025. All rights reserved.
 */
#ifndef WORKER_TOOL_H
#define WORKER_TOOL_H

#include "Logger.h"
#include "Segment.h"
#include "ThreadPool/ThreadPool.h"
#include <iostream>
#include <map>
#include <semaphore.h>

#ifndef IGNORE_PYBIND
#include "pybind11/pybind11.h"
#include "pybind11/stl.h"
#endif

#if defined(USE_NPU)
#include "acl/acl.h"
#endif

#define READ_LAYER_BUF_NUM          (2)    // TODO: 是KV的意思么？， Ping Pang Buffer

// TODO(yangxianpku: 2025.08.19): 是什么含义
#define WRITE_PREFILL_LAYER_BUF_NUM (1)   // Vllm v1
#define WRITE_DECODE_LAYER_BUF_NUM  (1)
#define LAYER_BUF_NUM                                            \
            (READ_LAYER_BUF_NUM + WRITE_PREFILL_LAYER_BUF_NUM +  \
            WRITE_DECODE_LAYER_BUF_NUM)


// TODO(yangxianpku: 2025.08.19): 一个Key或Value Buf固定256MB
#define SUB_BUF_SIZE (256 * 1024ULL * 1024ULL) // 每一层的Key or Value内存大小
#define LAYER_BUF_SIZE (SUB_BUF_SIZE * 2)      // 每一层的内存大小
#define TOTAL_LAYER_BUF_SIZE (LAYER_BUF_SIZE * LAYER_BUF_NUM)


/**
 *! @brief 定义Memory Copy方向
 * 
 */
// TODO(yangxianpku: 2025.08.21):  使用`class enum`形式, 移除typedef模式
typedef enum {
    DEVICE_TO_HOST,
    HOST_TO_DEVICE,
    DEVICE_TO_DEVICE,
    HOST_TO_HOST,
    BUTT
} DIRECTION;


// TODO(yangxianpku: 2025.08.19): 为什么不放在Logger.h中？
#define DEFAULT_LOG_FILE_SIZE (10 * 1024 * 1024)

typedef struct tag_BlkInfo {
    int64_t hashId;
    int32_t phyBlkId;
    Segment segment;
    tag_BlkInfo(int64_t hashId, int32_t phyBlkId, Segment segment)
        : hashId(hashId), phyBlkId(phyBlkId), segment(segment) {}
} BlkInfo;

/* =================Layer Buf Layout===========================
 * 0.......................................................1GB
 * |                          GPU 0                         |
 * 0..............................................512MB....1GB
 * |                    BUF 0                       | BUF 1 |
 * 0..................................256MB.......512MB
 * |             KEY BUF                | VALUE BUF |
 * 0.....128KB...256KB..256MB-128KB...256MB
 * | BLK 0 | BLK 1 | ....... | BLK 2047 |
 */

// TODO(yangxianpku: 2025.08.19): 这里怎么会用union呢?
typedef union {
#if defined(USE_GPU)
    cudaStream_t cudaStream;
#elif defined(USE_NPU)
    aclrtStream npuStream;
#endif
} DeviceStream;


/**
 *! @brief 每层的KV Cache信息
 * 
 */
// TODO(yangxianpku: 2025.08.21):  移除typedef模式
typedef struct {
    /* =====初始化时设置===== */
    uint8_t *    keyBuf;        // 单层Key Cache数据地址
    int64_t      keyBufLen;     // 单层Key Cache数据的大小(字节)
    uint8_t *    valueBuf;      // 单层Value Cache数据地址
    int64_t      valueBufLen;   // 单层Value Cache数据的大小(字节)
    DeviceType   type;          // Device类型: CPU, GPU, NPU
    bool         pinned_memory; // 是否为锁页内存
    DeviceStream stream;        // 设备工作流

    /* =====每次重排布时设置===== */
    std::vector<BlkInfo> blkInfo; // keyBuf和valueBuf中数据的元数据信息
    int32_t layerIdx;             // 需要重排布的层ID
    int32_t blockSize;            // 每个block中的token个数
    int32_t numHead;              // 每个token的head个数(全部，所有TP的头)
    int32_t headSize;             // 每个head的维度
    int32_t elementSize;          // 每个元素的字节数
    bool    hasValue;             // 标记是否有value，MLA无value(共享KV)
    int32_t tpRank; // 本worker的序号，TP并行时有效，取值为[0, tpNum)
    int32_t tpNum;  // TP并行数

    /* =====每次重排布时计算得到=====*/
    int32_t transferUnitSize;    // segment与layerBuf间最大的可连续拷贝的内存大小
    std::vector<int32_t> indexs; // 物理块连续情况
    int32_t tokenBytes;          // 单层Key/Value中，每个token的字节数
    int32_t blockBytes;          // 单层Key/Value中，每个block的字节数
    int32_t keyPos;              // 当前layer的key在segment内的偏移
    int32_t valuePos;            // 当前layer的value在segment内的偏移
    int32_t headOffset;          // TP并行时，head在token内部的起始偏移
} layerBuf;



/**
 * @brief 
 * 
 */
class WorkerTool {
private:
    DeviceType _deviceType;
    int32_t    _deviceId = -1;

    // 共3个buf，前2个用于读，第3个用于写
    layerBuf   _layerBufs[LAYER_BUF_NUM];   // 4

    ThreadPool _threadPool;

    // 用于read stream和exec stream间同步的信号量
    sem_t      _readSem;

    // 完成推理的层数，startLayerFill中初始化为0，waitLayerFillFinish中加1
    int32_t    _execLayerCnt;

    // 总层数
    int32_t    _totalLayerNum; 

    uint64_t   _totalLayerWriteTime = 0;


public:
    WorkerTool() : WorkerTool(DEVICE_TYPE_GPU) {}

    WorkerTool(DeviceType type) 
        : _threadPool(1), _deviceType(type) {
        const std::string LOG_PATH = \
            ConfigLoader::get_config<std::string>(
                ConfigLoader::EnvVars::LOG_PATH, 
                DEFAULT_LOG_PATH
            );
        
        Logger::instance().init(LOG_PATH + "WorkerTool_",
                                DEFAULT_LOG_FILE_SIZE);
        _threadPool.init();  // 线程池初始化但无任何任务

        initLayerBuf(false);
  
        /**
         *! @brief 初始化一个未命名的信号量
         * 
         * @param sem      要初始化的信号量对象的指针
         * @param pshared  信号量共享范围(0线程共享,1进程共享)
         * @param value    信号量的初始计数值， 通常为 ≥0 的整数：
         *                     - 0：初始无资源，获取信号量的线程会阻塞
         *                     - N：初始允许 N 个线程同时访问资源
         * @return int 
         */
        // int sem_init (sem_t *__sem, int __pshared, unsigned int __value);
        sem_init(&_readSem, 0, 0);
    }

    ~WorkerTool() {
        _threadPool.shutdown();

        sem_destroy(&_readSem);        // 清理信号量
        destroyLayerBuf();
        Logger::instance().shutdown();
    }

    /**
     *! @brief 获取指定指定索引的layerBuf
     * 
     * @param idx layerBuf索引, 取值[0, LAYER_BUF_NUM -1)
     * 
     * @return layerBuf& 
     */
    layerBuf &getLayerBuf(int32_t idx);



    /**
     *! @brief  将blockId的数据在DRAM和VRAM之间交换
     * 
     * @param segments      CPU DRAM中的segment对象
     * @param device_cache  每层KVCache指针地址，向量长度为模型层数
     * @param blockIds      需要读取的块ID，用于在dst中计算目标显存地址
     * @param direction     数据传输方向，CPU->GPU或GPU->CPU
     * @param num_layer     模型层数
     * @param num_blocks    vLLM引擎KV Cache表的总块数
     * @param block_size    每个block中的token个数
     * @param num_kv_heads  模型的总KV头数(总头数，不考虑PP)
     * @param head_size     每个head的维度(Q, KV的头一般都一样大)
     * @param element_size  每个KV Cache元素的字节数{fp16: 2, fp32:4, int8: 1, ...}
     * @param has_value     标记是否有value，MLA无value(MLA共享KV, 可认为只有K无V)
     * 
     * @return int32_t      执行成功返回0，失败返回非0
     */
    // int32_t swap(std::vector<Segment> &segments,
    //              std::vector<void *> const &device_cache,
    //              std::vector<int32_t> const &blockIds, 
    //              int32_t direction,
    //              int32_t num_layer, 
    //              int32_t num_blocks, 
    //              int32_t block_size,
    //              int32_t num_kv_heads, 
    //              int32_t head_size, 
    //              int32_t element_size,
    //              int32_t has_value);


    /**
     *! @brief 更新层的_layerBufs信息
     * 
     * @param readIdxSegmentMap   将读取的Segment字典
     * @param writeIdxSegmentMap  将写入的Segment字典
     * @param startLayer          模型起始层
     * @param block_size          vLLM Paged Attn块大小
     * @param numHead             模型注意力总头数(所有TP)
     * @param headSize            单个头的大小
     * @param elementSize         KV Cache单个元素的字节大小
     * @param hasValue            是否包含V Cache
     * @param tpRank              当前进程的张量并行rank
     * @param tpNum               模型推理部署的张量并行大小
     * 
     * @return int32_t 成功返回0否则返回具体错误码
     */
    int32_t updateLayerBufInfo(std::map<int32_t, Segment> &readIdxSegmentMap,
                               std::map<int32_t, Segment> &writeIdxSegmentMap,
                               int32_t startLayer, 
                               int32_t block_size,
                               int32_t numHead, 
                               int32_t headSize,
                               int32_t elementSize, 
                               bool    hasValue,
                               int32_t tpRank, 
                               int32_t tpNum);


    /**
     * @brief 
     * 
     * @param prefillReadIdxSegmentMap 
     * @param prefillWriteIdxSegmentMap 
     * @param decodeWriteIdxSegmentMap 
     * @param startLayer 
     * @param blockSize 
     * @param numHead 
     * @param headSize 
     * @param elementSize 
     * @param hasValue 
     * @param tpRank 
     * @param tpNum 
     * @return int32_t 
     */

    using SegmentMap = std::map<int32_t, Segment>;
    int32_t adjustUpdateLayerBufInfo(SegmentMap &prefillReadIdxSegmentMap,
                                    SegmentMap &prefillWriteIdxSegmentMap,
                                    SegmentMap &decodeWriteIdxSegmentMap,
                                    int32_t startLayer, 
                                    int32_t blockSize, 
                                    int32_t numHead,
                                    int32_t headSize, 
                                    int32_t elementSize, 
                                    bool    hasValue, 
                                    int32_t tpRank,
                                    int32_t tpNum
                                );


    /**
     *! @brief 将_layerBufs中的数据都Load到vLLM的KV Cache中
     * 
     * @param kvCachePtrs KV Cache Tensor的指针列表
     * @param numBlocks 
     * 
     * @return int32_t 成功返回0否则返回错误码
     */
    int32_t startLayerFill(std::vector<void *> &kvCachePtrs, int32_t numBlocks);


    /**
     *! @brief 将单层KV Cache数据下刷到Sgement中(Prefill阶段)
     * 
     * @param kvCachePtr 单层KV Cache的地址指针
     * @param numBlocks vLLM中Paged Attn的块数量
     * @param event 
     * 
     * @return int32_t 成功返回0否则返回错误状态码
     */
    int32_t singleLayerWrite(void *kvCachePtr, 
                            int32_t numBlocks, 
                            void *event
                        );


    /**
     *! @brief decode阶段凑满一个block然后整块下刷到Sgement中
     * 
     * @param kvCachePtrs 
     * @param numBlocks 
     * @return int32_t 
     */
    int32_t allLayerWrite(std::vector<void *> &kvCachePtrs, 
                        int32_t numBlocks);



    /**
     *! @brief 等待read stream完成
     * 
     */
    void waitLayerFillFinish();



    /**
     *! @brief Set the Device Id object
     * 
     * @param id 
     */
    void setDeviceId(int32_t id) { _deviceId = id; }


private:
    /**
     *! @brief 使用mmap申请内存空间，给key、value buf都指定地址
     * 
     * @param pinned_memory 
     * @return int32_t 
     */
    int32_t initLayerBuf(bool pinned_memory = false);


    /**
     *! @brief 释放层占用的内存空间
     * 
     * @return int32_t 
     */
    int32_t destroyLayerBuf();
};



typedef struct {
    int32_t layerIdx;
    int32_t numBlocks;
    int32_t transferUnitSize;
    int32_t hasValue;
    int32_t tpNum;
    int32_t tpRank;
    DeviceStream *stream;
} LayCacheInfo;



/**
 *! @brief  不开启TP的情况下，按层在DRAM和GPU之间拷贝数据
 * 
 * @param segments         CPU DRAM上的数据
 * @param layer_cache_ptr  Device上的KVCache指针
 * @param blockIds         Device端需要拷贝数据的物理块ID
 * @param direction        数据拷贝方向
 * @param type             Device类型
 * @param info             Device KV Cache辅助信息
 * @return int32_t 
 */
// TODO(yangxianpku: 2025.08.21): 这里不传引用？
int32_t swapOneLayerWithoutTP(std::vector<Segment> segments,
                            const void *layer_cache_ptr,
                            std::vector<int32_t> const &blockIds,
                            DIRECTION direction, 
                            DeviceType type,
                            LayCacheInfo &info
                        );

#ifndef IGNORE_PYBIND
namespace py = pybind11;

/**
 *! @brief 封装为Python接口
 * 
 */
PYBIND11_MODULE(WorkerTool, m) {
    m.doc() = "pybind11 WorkerTool";
    py::class_<WorkerTool>(m, "WorkerTool")
        .def(py::init<DeviceType>())
        // .def("swap", &WorkerTool::swap)
        .def("updateLayerBufInfo", &WorkerTool::updateLayerBufInfo)
        .def("adjustUpdateLayerBufInfo", &WorkerTool::adjustUpdateLayerBufInfo)
        .def("startLayerFill", &WorkerTool::startLayerFill)
        .def("singleLayerWrite", &WorkerTool::singleLayerWrite)
        .def("allLayerWrite", &WorkerTool::allLayerWrite)
        .def("waitLayerFillFinish", &WorkerTool::waitLayerFillFinish)
        .def("setDeviceId", &WorkerTool::setDeviceId);
}
#endif

#endif
